/* ****************************************************************** **
**    OpenSees - Open System for Earthquake Engineering Simulation    **
**          Pacific Earthquake Engineering Research Center            **
**                                                                    **
** ****************************************************************** */

//
// 文件: GPUMatrixAssembly.h
// 创建时间: 2025年10月2日
// 说明: GPU全局刚度矩阵组装核函数声明
// 作者: OpenSees GPU加速项目组
//
// 功能:
//   - 声明GPU并行矩阵组装核函数
//   - 支持原子操作组装（Phase 1）
//   - 支持单元着色优化组装（Phase 2，预留接口）
//   - CSR格式全局刚度矩阵组装
//

#ifndef GPUMatrixAssembly_h
#define GPUMatrixAssembly_h

#include <cuda_runtime.h>

//==============================================================================
// GPU矩阵组装核函数声明
//==============================================================================

/**
 * @brief GPU并行组装全局刚度矩阵（原子操作版本）
 *
 * 算法描述:
 * - 每个线程处理一个单元
 * - 遍历单元刚度矩阵的每个元素
 * - 根据DOF映射查找CSR格式中的对应位置
 * - 使用atomicAdd累加到全局刚度矩阵（避免竞争条件）
 *
 * 性能特点:
 * - 简单直接，易于实现和调试
 * - 原子操作有一定性能开销
 * - 适合作为基础实现和正确性验证
 * - Phase 2将用单元着色优化替代原子操作
 *
 * @param d_globalK 全局刚度矩阵值数组（CSR格式，GPU）
 * @param d_rowPtr CSR行指针数组（GPU）
 * @param d_colInd CSR列索引数组（GPU）
 * @param d_elementK 单元刚度矩阵缓冲区（GPU）
 * @param d_elementDOF 单元DOF映射数组（GPU）
 * @param numElements 单元数量
 * @param maxDOFPerElement 单元最大DOF数（例如：Quad=8, Beam3D=12）
 *
 * @note 要求d_globalK已经清零（在调用前使用cudaMemset）
 * @note 使用atomicAdd确保多个线程同时写入同一位置时的正确性
 * @note DOF索引为负数表示该DOF被约束，将被跳过
 */
__global__ void assembleGlobalMatrix_Atomic(
    double* d_globalK,
    int* d_rowPtr,
    int* d_colInd,
    double* d_elementK,
    int* d_elementDOF,
    int* d_elementDOFCount,
    int* d_feToElementMap,  // ✅ FE→Element索引映射
    int numElements,
    int maxFEDOF,           // ✅ FE DOF映射的布局参数
    int maxDOFPerElement    // Element刚度矩阵的布局参数
);

/**
 * @brief GPU并行组装全局刚度矩阵（单元着色优化版本）
 *
 * 算法描述:
 * - 预处理阶段对单元进行着色，保证同色单元不共享DOF
 * - 每次处理一个颜色的所有单元
 * - 同色单元并行组装时无需原子操作（无竞争）
 *
 * 性能优势:
 * - 消除原子操作开销
 * - 更高的内存访问效率
 * - 预期比原子操作版本快2-3倍
 *
 * @param d_globalK 全局刚度矩阵值数组（CSR格式，GPU）
 * @param d_rowPtr CSR行指针数组（GPU）
 * @param d_colInd CSR列索引数组（GPU）
 * @param d_elementK 单元刚度矩阵缓冲区（GPU）
 * @param d_elementDOF 单元DOF映射数组（GPU）
 * @param d_coloredElementIndices 当前颜色的单元索引列表（GPU）
 * @param numElementsInColor 当前颜色的单元数量
 * @param maxDOFPerElement 单元最大DOF数
 *
 * @note Phase 2实现（当前为预留接口）
 * @note 需要配合单元着色算法使用（GPUElementColoring）
 */
__global__ void assembleGlobalMatrix_ColorBatch(
    double* d_globalK,
    int* d_rowPtr,
    int* d_colInd,
    double* d_elementK,
    int* d_elementDOF,
    int* d_coloredElementIndices,
    int numElementsInColor,
    int maxDOFPerElement
);

//==============================================================================
// GPU辅助设备函数声明
//==============================================================================

/**
 * @brief GPU设备函数：在CSR格式中查找(row, col)位置
 *
 * @param row 行索引
 * @param col 列索引
 * @param d_rowPtr CSR行指针数组
 * @param d_colInd CSR列索引数组
 *
 * @return CSR值数组中的位置索引，如果未找到返回-1
 */
__device__ int findCSRPosition(
    int row,
    int col,
    const int* d_rowPtr,
    const int* d_colInd
);

//==============================================================================
// C接口声明（用于从C++代码调用）
//==============================================================================

#ifdef __cplusplus
extern "C" {
#endif

/**
 * @brief 启动原子操作矩阵组装核函数
 *
 * @param numElements 单元数量
 * @param d_globalK 全局刚度矩阵（GPU）
 * @param d_rowPtr CSR行指针（GPU）
 * @param d_colInd CSR列索引（GPU）
 * @param d_elementK 单元刚度矩阵（GPU）
 * @param d_elementDOF 单元DOF映射（GPU）
 * @param maxDOFPerElement 单元最大DOF数
 * @param stream CUDA流（用于异步执行）
 *
 * @return 0-成功，-1-失败
 */
int launchAtomicAssemblyKernel(
    int numElements,
    double* d_globalK,
    const int* d_rowPtr,
    const int* d_colInd,
    const double* d_elementK,
    const int* d_elementDOF,
    const int* d_elementDOFCount,
    const int* d_feToElementMap,  // ✅ FE→Element索引映射
    int maxFEDOF,                  // ✅ FE DOF映射布局
    int maxDOFPerElement,
    cudaStream_t stream = 0
);

/**
 * @brief 启动单元着色矩阵组装核函数
 *
 * @param numElementsInColor 当前颜色的单元数量
 * @param d_globalK 全局刚度矩阵（GPU）
 * @param d_rowPtr CSR行指针（GPU）
 * @param d_colInd CSR列索引（GPU）
 * @param d_elementK 单元刚度矩阵（GPU）
 * @param d_elementDOF 单元DOF映射（GPU）
 * @param d_coloredElementIndices 当前颜色的单元索引（GPU）
 * @param maxDOFPerElement 单元最大DOF数
 * @param stream CUDA流（用于异步执行）
 *
 * @return 0-成功，-1-失败
 *
 * @note Phase 2实现
 */
int launchColoredAssemblyKernel(
    int numElementsInColor,
    double* d_globalK,
    const int* d_rowPtr,
    const int* d_colInd,
    const double* d_elementK,
    const int* d_elementDOF,
    const int* d_coloredElementIndices,
    int maxDOFPerElement,
    cudaStream_t stream = 0
);

/**
 * @brief 包装函数：用于formTangent_GPU()调用GPU并行组装
 */
int launchAssemblyKernel_Atomic(
    double* d_globalK,
    int* d_rowPtr,
    int* d_colInd,
    double* d_feStiffness,
    int* d_feDOFMapping,
    int* d_feDOFCount,
    int* d_feToElementMap,  // ✅ FE→Element索引映射
    int numFE,
    int maxFEDOF,            // ✅ FE DOF映射的布局参数
    int maxDOFPerElement,    // Element刚度矩阵的布局参数
    int gridSize,
    int blockSize
);

#ifdef __cplusplus
}
#endif

#endif // GPUMatrixAssembly_h
